#!/usr/bin/env python
# coding: utf-8

# In[10]:


import pandas as pd
import numpy as np
from copy import deepcopy
import warnings

class KLineProcessor:
    def __init__(self, df):
        self._validate_input(df)
        self.data = df

    @staticmethod
    def _validate_input(df: pd.DataFrame):
        required_cols = {'trade_date', 'high', 'low'}
        if not required_cols.issubset(df.columns):
            missing = required_cols - set(df.columns)
            raise ValueError(f"输入数据缺少必要列：{missing}")

    def _preprocess_data(self):
        self.data['Fmark'] = np.zeros(self.data.shape[0])
        self.data['Fval'] = np.zeros(self.data.shape[0])
        kline = [row.to_dict() for _, row in self.data.iterrows()]

        res = []
        for k in kline:
            k['Fmark'], k['Fval'], k['line'] = 0, None, None
            res.append(k)
        return res

    def _remove_merged_kline(self):
        kline = self.kline
        new_kline = kline[:2]

        # 初始化dir，dir表示目前是增长还是下降
        # dir ： 1上升， -1下降
        for k in kline[2:]:
            k1, k2 = new_kline[-2 : ]
            if k2['high'] > k1['high']:
                dir = 1
            elif k2['low'] < k1['low']:
                dir = -1
            else:
                dir = 0

            cur_high, cur_low = k['high'], k['low']
            last_high, last_low = k2['high'], k2['low']
            """
            如果上升
            高点与低点分别为两者高点与低点的较大值
            反之较小值
            """
            if (cur_high <= last_high and cur_low >= last_low) or (cur_high >= last_high and cur_low <= last_low):
                if dir == 1:
                   new_high = max(last_high, cur_high)
                   new_low = max(last_low, cur_low)
                elif dir == -1:
                    new_high = min(last_high, cur_high)
                    new_low = min(last_low, cur_low)
                else:
                    raise ValueError
                k['high'] = new_high
                k['low'] = new_low
                new_kline.pop(-1)
            new_kline.append(k)

        return new_kline

    def _identify_fractals(self):
        kn = self.processed_kline
        df1 = pd.DataFrame(kn)
        i = 0
        while(i < len(kn)):
            if i == 0 or i == len(kn) - 1:
                i += 1
                continue
            k1, k2, k3 = kn[i - 1: i + 2]
            i += 1
            if k1['high'] < k2['high'] > k3['high']:
                k2['Fmark'] = 1
                k2['Fval'] = k2['high']
            if k1['low'] > k2['low'] < k3['low']:
                k2['Fmark'] = -1
                k2['Fval'] = k2['low']
        self.processed_kline = kn
        # df = pd.DataFrame(fractals)
        # df.to_csv('1.csv', index=False)
        fractals = [{"trade_date" : x['trade_date'], "Fmark" : x['Fmark'], "Fval" : x['Fval']} for x in self.processed_kline if x['Fmark'] in [1, -1]]

        return fractals

    def process_kline_fractals(self):
        kline = self.kline
        new_kline = kline[:2]
        fractals = []

        for i, k in enumerate(kline[2:], start=2):
            k1, k2 = new_kline[-2:]
            if k2['high'] > k1['high']:
                direction = 1
            elif k2['low'] < k1['low']:
                direction = -1
            else:
                direction = 1

            cur_high, cur_low = k['high'], k['low']
            last_high, last_low = k2['high'], k2['low']

            if (cur_high <= last_high and cur_low >= last_low) or (cur_high >= last_high and cur_low <= last_low):
                if direction == 1:
                    new_high = max(last_high, cur_high)
                    new_low = max(last_low, cur_low)
                elif direction == -1:
                    new_high = min(last_high, cur_high)
                    new_low = min(last_low, cur_low)
                else:
                    raise ValueError
                k['high'] = new_high
                k['low'] = new_low
                new_kline.pop(-1)

            new_kline.append(k)

            if i >= 2 and i <= len(kline) - 1 and len(new_kline) >= 3:
                k1, k2, k3 = new_kline[-3:]
                if k1['high'] < k2['high'] > k3['high']:
                    if (not fractals or fractals[-1]["trade_date"] != k2['trade_date']):
                        k2['Fmark'] = 1
                        k2['Fval'] = k2['high']
                        fractals.append({"trade_date": k2['trade_date'], "Fmark": k2['Fmark'], "Fval": k2['Fval']})
                if k1['low'] > k2['low'] < k3['low']:
                    if not fractals or fractals[-1]["trade_date"] != k2['trade_date']:
                        k2['Fmark'] = -1
                        k2['Fval'] = k2['low']
                        fractals.append({"trade_date": k2['trade_date'], "Fmark": k2['Fmark'], "Fval": k2['Fval']})

        self.processed_kline = new_kline

        return fractals

    def getpreLine(self):
        kn = self.processed_kline
        Fp = sorted(deepcopy(self.fractals), key=lambda x: x['trade_date'], reverse=False)
        line = []
        # kn_df = pd.DataFrame(kn)
        # kn_df.to_csv('example1.csv', index=False)  # index=False 表示不保存行索引
        """
        当前的分型相较于栈顶的分型类型相同，如果当前更高或者更低将栈弹出，以当前作为新的笔结束点。
                                 如果不同，那么作为笔的结束，需要满足与之前分型中间有一个k线的条件
        """
        for i in range(len(Fp)):
            k = Fp[i]
            if len(line) == 0:
                line.append(k)
            else:
                kpre = line[-1]
                # print(kpre, k)
                if kpre['Fmark'] == k['Fmark']:
                    if (kpre['Fmark'] == 1 and kpre['Fval'] < k['Fval']) or (
                            kpre['Fmark'] == -1 and kpre['Fval'] > k['Fval']):
                        line.pop(-1)
                        line.append(k)
                else:
                    # 如果前面一个是顶分型那么下一个分型必须是底分型，这样下一个分型的值必须小于上一个分型的值，否则上一个顶分型会被pop掉。反之同理。
                    if (kpre['Fmark'] == 1 and k['Fval'] >= kpre['Fval']) or (
                            kpre['Fmark'] == -1 and k['Fval'] <= kpre['Fval']):
                        line.pop(-1)
                        continue

                    knum = [x for x in kn if kpre['trade_date'] <= x['trade_date'] <= k['trade_date']]
                    if len(knum) >= 5:
                        line.append(k)
                        # 当确定了一个笔后如果新加入的这个笔为顶，那么如果与上一个笔之间有更高的点，那么这个笔会被去除，底同理。
                        maxx, minn = 0, 0x3f3f3f3f
                        for x in kn:
                            if kpre['trade_date'] <= x['trade_date'] <= k['trade_date']:
                                maxx = max(x['high'], maxx)
                                minn = min(x['low'], minn)
                        if (kpre['Fmark'] == -1 and k['Fval'] < maxx) or (kpre['Fmark'] == 1 and k['Fval'] > minn):
                            line.pop(-1)
        return line

    def getpreLine_plus(self):
        kn = self.processed_kline
        Fp = sorted(self.fractals, key=lambda  x:x['trade_date'], reverse=False)
        line = []
        kn_df = pd.DataFrame(kn)
        # Fp_df = pd.DataFrame(Fp)
        kn_df.to_csv('data1.csv', index=False)  # index=False 表示不保存行索引
        """
        当前的分型相较于栈顶的分型类型相同，如果当前更高或者更低将栈弹出，以当前作为新的笔结束点。
                                 如果不同，那么作为笔的结束，需要满足与之前分型中间有一个k线的条件
        """
        for i in range(len(Fp)):
            k = Fp[i]
            if len(line) <= 1:
                line.append(k)
            else :
                kpre = line[-1]
                kpre_plus = line[-2]

                if kpre['Fmark'] == k['Fmark']:
                    if (kpre['Fmark'] == 1 and kpre['Fval'] < k['Fval']) or (kpre['Fmark'] == -1 and kpre['Fval'] > k['Fval']):
                        line.pop(-1)
                        line.append(k)
                else:
                    if (kpre_plus['Fmark'] == 1 and kpre_plus['Fval'] < k['Fval']) or (
                            kpre_plus['Fmark'] == -1 and kpre_plus['Fval'] > k['Fval']):
                        line.pop(-2)
                        line.append(k)
                    # 如果前面一个是顶分型那么下一个分型必须是底分型，这样下一个分型的值必须小于上一个分型的值，否则上一个顶分型会被pop掉。反之同理。
                    if (kpre['Fmark'] == 1 and k['Fval'] >= kpre['Fval']) or (kpre['Fmark'] == -1 and k['Fval'] <= kpre['Fval']):
                        line.pop(-1)
                        continue


               # print(kpre, k)
                if kpre['Fmark'] == k['Fmark']:
                    if (kpre['Fmark'] == 1 and kpre['Fval'] < k['Fval']) or (kpre['Fmark'] == -1 and kpre['Fval'] > k['Fval']):
                        line.pop(-1)
                        line.append(k)
                else:


                    # 笔结束
                    knum = [x for x in kn if kpre['trade_date'] <= x['trade_date'] <= k['trade_date']]
                    if len(knum) >= 5:
                        line.append(k)
                        # 当确定了一个笔后如果新加入的这个笔为顶，那么如果与上一个笔之间有更高的点，那么这个笔会被去除，底同理。
                        maxx, minn = 0, 0x3f3f3f3f
                        condition = (kn_df['trade_date'] >= kpre['trade_date']) & (kn_df['trade_date'] <= k['trade_date'])
                        true_indexes = condition[condition].index
                        for i in true_indexes:
                            maxx = max(kn[i]['high'], maxx)
                            minn = min(kn[i]['low'], minn)
                        if (kpre['Fmark'] == -1 and k['Fval'] < maxx) or (kpre['Fmark'] == 1 and k['Fval'] > minn):
                            line.pop(-1)
        return line

    def getLine(self):
        line = self.getpreLine()

        datelist = [x['trade_date'] for x in line]
        for k in self.processed_kline:
            if k['trade_date'] in datelist:
                k['line'] = k['Fval']
        return line

    def get_fractals(self):
        flag = 0
        for i in range(self.fractals.shape[0]):
            if self.fractals[i][-2] != 0:
                flag = self.fractals[i][-2]  # 倒数第二列为 Fmark={1,-1,0}
                continue
            if flag == -1:
                self.fractals[i][-2] = 2
            if flag == 1:
                self.fractals[i][-2] = -2
        colum_idx = self.data.columns.get_loc('trade_date')
        self.fractals[0: int(np.where(self.fractals[:, colum_idx] == self.line['trade_date'][0])[0]), -2] = -2 if self.line['Fmark'].iloc[0] == -1 else 2
        return self.fractals

    def get_data(self):
        self.kline = self._preprocess_data()
        self.fractals = self.process_kline_fractals()
        self.line = self.getLine()
        self.line = pd.DataFrame(self.line)

        # 将笔的数据，加入到raw_data中
        for j in range(self.line.shape[0]):
            self.data.loc[self.data['trade_date'] == self.line['trade_date'].iloc[j], 'Fmark'] = self.line['Fmark'].iloc[j]
            self.data.loc[self.data['trade_date'] == self.line['trade_date'].iloc[j], 'Fval'] = self.line['Fval'].iloc[j]

        self.fractals = self.data.values
        self.fractals = self.get_fractals()
        self.data = pd.DataFrame(self.fractals, columns=self.data.columns)

        # 定义替换规则
        replace_dict = {1: 0, -1: 1, -2: 3}
        # 替换指定列中的值
        self.data['Fmark'] = self.data['Fmark'].replace(replace_dict)

        # 数据类型映射
        dtype_mapping = {
            'trade_date': 'datetime64[ns]',  # 日期列
            'open': 'float32',  # 开盘价
            'high': 'float32',  # 最高价
            'low': 'float32',  # 最低价
            'close': 'float32',  # 价格列
            'Fmark': 'int',  # 标签
            'Fval': 'float32',
        }
        self.data = self.data.astype(dtype_mapping)

        return self.data

if __name__ == '__main__':
    df = pd.read_csv('data1.csv')
    f = pd.DataFrame({'trade_date': [1, 2, 3], 'Fmark': [1, -1, -2]})
    df = df.sort_values('trade_date', ascending=True).reset_index(drop=True)
    warnings.filterwarnings('error', category=pd.errors.PerformanceWarning)
    L = KLineProcessor(df)
    df1 = L.get_data()
    print(df1[["trade_date", "Fmark"]])


# In[11]:


import pandas as pd
import numpy as np
from copy import deepcopy
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


class KLineProcessor:
    def __init__(self, df):
        self._validate_input(df)
        self.data = df

    @staticmethod
    def _validate_input(df: pd.DataFrame):
        required_cols = {'trade_date', 'high', 'low'}
        if not required_cols.issubset(df.columns):
            missing = required_cols - set(df.columns)
            raise ValueError(f"输入数据缺少必要列：{missing}")

    def _preprocess_data(self):
        self.data['Fmark'] = np.zeros(self.data.shape[0])
        self.data['Fval'] = np.zeros(self.data.shape[0])
        kline = [row.to_dict() for _, row in self.data.iterrows()]

        res = []
        for k in kline:
            k['Fmark'], k['Fval'], k['line'] = 0, None, None
            res.append(k)
        return res

    def _remove_merged_kline(self):
        kline = self.kline
        new_kline = kline[:2]

        # 初始化dir，dir表示目前是增长还是下降
        # dir ： 1上升， -1下降
        for k in kline[2:]:
            k1, k2 = new_kline[-2:]
            if k2['high'] > k1['high']:
                dir = 1
            elif k2['low'] < k1['low']:
                dir = -1
            else:
                dir = 0

            cur_high, cur_low = k['high'], k['low']
            last_high, last_low = k2['high'], k2['low']
            """
            如果上升
            高点与低点分别为两者高点与低点的较大值
            反之较小值
            """
            if (cur_high <= last_high and cur_low >= last_low) or (cur_high >= last_high and cur_low <= last_low):
                if dir == 1:
                    new_high = max(last_high, cur_high)
                    new_low = max(last_low, cur_low)
                elif dir == -1:
                    new_high = min(last_high, cur_high)
                    new_low = min(last_low, cur_low)
                else:
                    raise ValueError
                k['high'] = new_high
                k['low'] = new_low
                new_kline.pop(-1)
            new_kline.append(k)

        return new_kline

    def _identify_fractals(self):
        kn = self.processed_kline
        df1 = pd.DataFrame(kn)
        i = 0
        while (i < len(kn)):
            if i == 0 or i == len(kn) - 1:
                i += 1
                continue
            k1, k2, k3 = kn[i - 1: i + 2]
            i += 1
            if k1['high'] < k2['high'] > k3['high']:
                k2['Fmark'] = 1
                k2['Fval'] = k2['high']
            if k1['low'] > k2['low'] < k3['low']:
                k2['Fmark'] = -1
                k2['Fval'] = k2['low']
        self.processed_kline = kn
        fractals = [{"trade_date": x['trade_date'], "Fmark": x['Fmark'], "Fval": x['Fval']} for x in self.processed_kline if
                    x['Fmark'] in [1, -1]]

        return fractals

    def process_kline_fractals(self):
        kline = self.kline
        new_kline = kline[:2]
        fractals = []

        for i, k in enumerate(kline[2:], start=2):
            k1, k2 = new_kline[-2:]
            if k2['high'] > k1['high']:
                direction = 1
            elif k2['low'] < k1['low']:
                direction = -1
            else:
                direction = 1

            cur_high, cur_low = k['high'], k['low']
            last_high, last_low = k2['high'], k2['low']

            if (cur_high <= last_high and cur_low >= last_low) or (cur_high >= last_high and cur_low <= last_low):
                if direction == 1:
                    new_high = max(last_high, cur_high)
                    new_low = max(last_low, cur_low)
                elif direction == -1:
                    new_high = min(last_high, cur_high)
                    new_low = min(last_low, cur_low)
                else:
                    raise ValueError
                k['high'] = new_high
                k['low'] = new_low
                new_kline.pop(-1)

            new_kline.append(k)

            if i >= 2 and i <= len(kline) - 1 and len(new_kline) >= 3:
                k1, k2, k3 = new_kline[-3:]
                if k1['high'] < k2['high'] > k3['high']:
                    if (not fractals or fractals[-1]["trade_date"] != k2['trade_date']):
                        k2['Fmark'] = 1
                        k2['Fval'] = k2['high']
                        fractals.append({"trade_date": k2['trade_date'], "Fmark": k2['Fmark'], "Fval": k2['Fval']})
                if k1['low'] > k2['low'] < k3['low']:
                    if not fractals or fractals[-1]["trade_date"] != k2['trade_date']:
                        k2['Fmark'] = -1
                        k2['Fval'] = k2['low']
                        fractals.append({"trade_date": k2['trade_date'], "Fmark": k2['Fmark'], "Fval": k2['Fval']})

        self.processed_kline = new_kline

        return fractals

    def getpreLine(self):
        kn = self.processed_kline
        Fp = sorted(deepcopy(self.fractals), key=lambda x: x['trade_date'], reverse=False)
        line = []
        for i in range(len(Fp)):
            k = Fp[i]
            if len(line) == 0:
                line.append(k)
            else:
                kpre = line[-1]
                if kpre['Fmark'] == k['Fmark']:
                    if (kpre['Fmark'] == 1 and kpre['Fval'] < k['Fval']) or (
                            kpre['Fmark'] == -1 and kpre['Fval'] > k['Fval']):
                        line.pop(-1)
                        line.append(k)
                else:
                    if (kpre['Fmark'] == 1 and k['Fval'] >= kpre['Fval']) or (
                            kpre['Fmark'] == -1 and k['Fval'] <= kpre['Fval']):
                        line.pop(-1)
                        continue

                    knum = [x for x in kn if kpre['trade_date'] <= x['trade_date'] <= k['trade_date']]
                    if len(knum) >= 5:
                        line.append(k)
                        maxx, minn = 0, 0x3f3f3f3f
                        for x in kn:
                            if kpre['trade_date'] <= x['trade_date'] <= k['trade_date']:
                                maxx = max(x['high'], maxx)
                                minn = min(x['low'], minn)
                        if (kpre['Fmark'] == -1 and k['Fval'] < maxx) or (kpre['Fmark'] == 1 and k['Fval'] > minn):
                            line.pop(-1)
        return line

    def getpreLine_plus(self):
        kn = self.processed_kline
        Fp = sorted(self.fractals, key=lambda x: x['trade_date'], reverse=False)
        line = []
        for i in range(len(Fp)):
            k = Fp[i]
            if len(line) <= 1:
                line.append(k)
            else:
                kpre = line[-1]
                kpre_plus = line[-2]

                if kpre['Fmark'] == k['Fmark']:
                    if (kpre['Fmark'] == 1 and kpre['Fval'] < k['Fval']) or (
                            kpre['Fmark'] == -1 and kpre['Fval'] > k['Fval']):
                        line.pop(-1)
                        line.append(k)
                else:
                    if (kpre_plus['Fmark'] == 1 and kpre_plus['Fval'] < k['Fval']) or (
                            kpre_plus['Fmark'] == -1 and kpre_plus['Fval'] > k['Fval']):
                        line.pop(-2)
                        line.append(k)
                    if (kpre['Fmark'] == 1 and k['Fval'] >= kpre['Fval']) or (
                            kpre['Fmark'] == -1 and k['Fval'] <= kpre['Fval']):
                        line.pop(-1)
                        continue

                if kpre['Fmark'] == k['Fmark']:
                    if (kpre['Fmark'] == 1 and kpre['Fval'] < k['Fval']) or (
                            kpre['Fmark'] == -1 and kpre['Fval'] > k['Fval']):
                        line.pop(-1)
                        line.append(k)
                else:
                    knum = [x for x in kn if kpre['trade_date'] <= x['trade_date'] <= k['trade_date']]
                    if len(knum) >= 5:
                        line.append(k)
                        maxx, minn = 0, 0x3f3f3f3f
                        condition = (pd.DataFrame(kn)['trade_date'] >= kpre['trade_date']) & (
                                pd.DataFrame(kn)['trade_date'] <= k['trade_date'])
                        true_indexes = condition[condition].index
                        for i in true_indexes:
                            maxx = max(kn[i]['high'], maxx)
                            minn = min(kn[i]['low'], minn)
                        if (kpre['Fmark'] == -1 and k['Fval'] < maxx) or (kpre['Fmark'] == 1 and k['Fval'] > minn):
                            line.pop(-1)
        return line

    def getLine(self):
        line = self.getpreLine()

        datelist = [x['trade_date'] for x in line]
        for k in self.processed_kline:
            if k['trade_date'] in datelist:
                k['line'] = k['Fval']
        return line

    def get_fractals(self):
        flag = 0
        for i in range(self.fractals.shape[0]):
            if self.fractals[i][-2] != 0:
                flag = self.fractals[i][-2]
                continue
            if flag == -1:
                self.fractals[i][-2] = 2
            if flag == 1:
                self.fractals[i][-2] = -2
        colum_idx = self.data.columns.get_loc('trade_date')
        self.fractals[0: int(np.where(self.fractals[:, colum_idx] == self.line['trade_date'][0])[0]), -2] = -2 if             self.line['Fmark'].iloc[0] == -1 else 2
        return self.fractals

    def get_data(self):
        self.kline = self._preprocess_data()
        self.fractals = self.process_kline_fractals()
        self.line = self.getLine()
        self.line = pd.DataFrame(self.line)

        for j in range(self.line.shape[0]):
            self.data.loc[self.data['trade_date'] == self.line['trade_date'].iloc[j], 'Fmark'] = self.line['Fmark'].iloc[
                j]
            self.data.loc[self.data['trade_date'] == self.line['trade_date'].iloc[j], 'Fval'] = self.line['Fval'].iloc[
                j]

        self.fractals = self.data.values
        self.fractals = self.get_fractals()
        self.data = pd.DataFrame(self.fractals, columns=self.data.columns)

        replace_dict = {1: 0, -1: 1, -2: 3}
        self.data['Fmark'] = self.data['Fmark'].replace(replace_dict)

        dtype_mapping = {
            'trade_date': 'datetime64[ns]',
            'open': 'float32',
            'high': 'float32',
            'low': 'float32',
            'close': 'float32',
            'Fmark': 'int',
            'Fval': 'float32',
        }
        self.data = self.data.astype(dtype_mapping)

        return self.data


def visualize_kline(data, fractals, lines):
    fig, ax = plt.subplots(figsize=(12, 6))

    # 绘制K线
    ax.vlines(data['trade_date'], data['low'], data['high'], color='black', linewidth=1)

    # 标记分型
    for fractal in fractals:
        trade_date = fractal['trade_date']
        fmark = fractal['Fmark']
        fval = fractal['Fval']
        if fmark == 1:
            ax.scatter(trade_date, fval, color='red', marker='^', s=50)
        elif fmark == -1:
            ax.scatter(trade_date, fval, color='green', marker='v', s=50)

    # 绘制笔
    for i in range(len(lines) - 1):
        start_date = lines[i]['trade_date']
        start_val = lines[i]['Fval']
        end_date = lines[i + 1]['trade_date']
        end_val = lines[i + 1]['Fval']
        ax.plot([start_date, end_date], [start_val, end_val], color='blue', linestyle='-', linewidth=2)

    # 设置日期格式
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    plt.xticks(rotation=45)

    plt.title('K线图及分型、笔标记')
    plt.xlabel('日期')
    plt.ylabel('价格')
    plt.grid(True)
    plt.show()


if __name__ == '__main__':
    df = pd.read_csv('data1.csv')
    df = df.sort_values('trade_date', ascending=True).reset_index(drop=True)
    L = KLineProcessor(df)
    df1 = L.get_data()

    fractals = L.fractals
    lines = L.line.to_dict(orient='records')

    visualize_kline(df1, fractals, lines)

    


# In[ ]:




